use prost::Message;

#[derive(Clone, PartialEq, Message)]
pub struct VectorProto {
    #[prost(string, tag="1")]
    pub id: String,
    #[prost(float, repeated, tag="2")]
    pub embedding: Vec<f32>,
    #[prost(string, tag="3")]
    pub metadata: String,
}

impl VectorProto {
    pub fn set_id(&mut self, id: String) {
        self.id = id;
    }

    pub fn get_id(&self) -> &str {
        &self.id
    }

    pub fn set_embedding(&mut self, embedding: Vec<f32>) {
        self.embedding = embedding;
    }

    pub fn get_embedding(&self) -> &[f32] {
        &self.embedding
    }

    pub fn set_metadata(&mut self, metadata: String) {
        self.metadata = metadata;
    }

    pub fn get_metadata(&self) -> &str {
        &self.metadata
    }
}

#[derive(Clone, PartialEq, Message)]
pub struct CollectionProto {
    #[prost(message, repeated, tag="1")]
    pub vectors: Vec<VectorProto>,
}

impl CollectionProto {
    pub fn mut_vectors(&mut self) -> &mut Vec<VectorProto> {
        &mut self.vectors
    }

    pub fn take_vectors(&mut self) -> Vec<VectorProto> {
        std::mem::take(&mut self.vectors)
    }

    pub fn write_to_vec(&self, buf: &mut Vec<u8>) -> Result<(), prost::EncodeError> {
        Message::encode(self, buf)
    }
}